# coding = utf-8

'''
HDenseunet V2
'''

import torch.nn as nn
from torchvision import models
import torch
import torch.nn.functional as F
from collections import OrderedDict

class DenseUNet(nn.Module):
    def __init__(self, in_ch=3, out_ch=3):
        super(DenseUNet, self).__init__()
        densenet = models.densenet161(pretrained=True)
        backbone = list(list(densenet.children())[0].children())

        if in_ch != 3:
            backbone[0] = nn.Conv2d(in_ch, 96, kernel_size=7, stride=2, padding=3, bias=False)

        self.conv1 = nn.Sequential(*backbone[:3])
        self.mp = backbone[3]
        self.denseblock1 = backbone[4]
        self.transition1 = backbone[5]
        self.denseblock2 = backbone[6]
        self.transition2 = backbone[7]
        self.denseblock3 = backbone[8]
        self.transition3 = backbone[9]
        self.denseblock4 = backbone[10]
        self.bn = backbone[11]
        self.up1 = _Up(x1_ch=2208, x2_ch=2112, out_ch=768)
        self.up2 = _Up(x1_ch=768, x2_ch=768, out_ch=384)
        self.up3 = _Up(x1_ch=384, x2_ch=384, out_ch=96)
        self.up4 = _Up(x1_ch=96, x2_ch=96, out_ch=96)
        self.up5 = nn.Sequential(
            _Interpolate(),
            nn.BatchNorm2d(num_features=96),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=96, out_channels=64, kernel_size=3, padding=1)
        )
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=out_ch, kernel_size=1)

        self.up1_conv = nn.Conv2d(in_channels=768, out_channels=out_ch, kernel_size=1)
        self.up2_conv = nn.Conv2d(in_channels=384, out_channels=out_ch, kernel_size=1)
        self.up3_conv = nn.Conv2d(in_channels=96, out_channels=out_ch, kernel_size=1)
        self.up4_conv = nn.Conv2d(in_channels=96, out_channels=out_ch, kernel_size=1)

    def forward(self, x):
        x = self.conv1(x)
        x_ = self.mp(x)
        x1 = self.denseblock1(x_)
        x1t = self.transition1(x1)
        x2 = self.denseblock2(x1t)
        x2t = self.transition2(x2)
        x3 = self.denseblock3(x2t)
        x3t = self.transition3(x3)
        x4 = self.denseblock4(x3t)
        x4 = self.bn(x4)
        x5 = self.up1(x4, x3)
        x6 = self.up2(x5, x2)
        x7 = self.up3(x6, x1)
        x8 = self.up4(x7, x)
        feat = self.up5(x8)
        cls = self.conv2(feat)

        up1_cls = self.up1_conv(x5)
        up2_cls = self.up2_conv(x6)
        up3_cls = self.up3_conv(x7)
        up4_cls = self.up4_conv(x8)

        return {'output': cls, 'up1_cls': up1_cls, 'up2_cls': up2_cls, 'up3_cls': up3_cls, 'up4_cls': up4_cls, "feature": feat}


class _Interpolate(nn.Module):
    def __init__(self, scale_factor=2, mode='bilinear', align_corners=True):
        super(_Interpolate, self).__init__()
        self.scale_factor = scale_factor
        self.mode = mode
        self.align_corners = align_corners

    def forward(self, x):
        x = nn.functional.interpolate(x, scale_factor=self.scale_factor, mode=self.mode,
                                      align_corners=self.align_corners)
        return x


class _Up(nn.Module):
    def __init__(self, x1_ch, x2_ch, out_ch):
        super(_Up, self).__init__()
        self.up = _Interpolate()
        self.conv1x1 = nn.Conv2d(in_channels=x2_ch, out_channels=x1_ch, kernel_size=1)
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=x1_ch, out_channels=out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x1, x2):
        x1 = self.up(x1)
        x2 = self.conv1x1(x2)
        x = x1 + x2
        x = self.conv(x)
        return x


# UNET 3D
class pub(nn.Module):

    def __init__(self, in_channels, out_channels, batch_norm=True):
        super(pub, self).__init__()
        inter_channels = in_channels/2 if in_channels > out_channels else out_channels/2
        inter_channels = int(inter_channels)
        layers = [
                    nn.Conv3d(in_channels, inter_channels, 3, stride=1, padding=1),
                    nn.ReLU(True),
                    nn.Conv3d(inter_channels, out_channels, 3, stride=1, padding=1),
                    nn.ReLU(True)
                 ]
        if batch_norm:
            layers.insert(1, nn.BatchNorm3d(inter_channels))
            layers.insert(len(layers)-1, nn.BatchNorm3d(out_channels))
        self.pub = nn.Sequential(*layers)

    def forward(self, x):
        return self.pub(x)


class unet3dDown(nn.Module):

    def __init__(self, in_channels, out_channels, batch_norm=True):
        super(unet3dDown, self).__init__()
        self.pub = pub(in_channels, out_channels, batch_norm)
        self.pool = nn.MaxPool3d(2, stride=2)

    def forward(self, x):
        x = self.pool(x)
        x = self.pub(x)
        return x


class unet3dUp(nn.Module):
    def __init__(self, in_channels, out_channels, batch_norm=True, sample=True):
        super(unet3dUp, self).__init__()
        self.pub = pub(int(in_channels/2+in_channels), out_channels, batch_norm)
        if sample:
            self.sample = nn.Upsample(scale_factor=2, mode='nearest')
        else:
            self.sample = nn.ConvTranspose3d(in_channels, in_channels, 2, stride=2)

    def forward(self, x, x1):
        x = self.sample(x)
        #c1 = (x1.size(2) - x.size(2)) // 2
        #c2 = (x1.size(3) - x.size(3)) // 2
        #x1 = x1[:, :, c1:-c1, c2:-c2, c2:-c2]
        x = torch.cat((x, x1), dim=1)
        x = self.pub(x)
        return x


class UNET3D(nn.Module):
    def __init__(self, init_channels=1, class_nums=1, batch_norm=True, sample=True):
        super(UNET3D, self).__init__()
        self.down1 = pub(init_channels, 64, batch_norm)
        self.down2 = unet3dDown(64, 128, batch_norm)
        self.down3 = unet3dDown(128, 256, batch_norm)
        self.down4 = unet3dDown(256, 512, batch_norm)
        self.up3 = unet3dUp(512, 256, batch_norm, sample)
        self.up2 = unet3dUp(256, 128, batch_norm, sample)
        self.up1 = unet3dUp(128, 64, batch_norm, sample)
        self.con_last = nn.Conv3d(64, class_nums, 1)
        #self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x1 = self.down1(x)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x = self.up3(x4, x3)
        x = self.up2(x, x2)
        last = self.up1(x, x1)
        x = self.con_last(last)
        return (last,x)

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv3d):
                nn.init.kaiming_uniform(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm3d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

#为了训练模型设计
class denseUnet3dForTrain(nn.Module):
    def __init__(self, drop_rate=0.3):
        super(denseUnet3dForTrain, self).__init__()
        self.dense_unet_3d = UNET3D(init_channels=1, class_nums=3)
        self.drop = drop_rate
        self.finalConv3d1 = nn.Conv3d(64, 64, (3, 3, 3), padding=(1, 1, 1))
        self.finalBn = nn.BatchNorm3d(64)
        self.finalAc = nn.ReLU(inplace=True)
        self.finalConv3d2 = nn.Conv3d(64, 3, (1, 1, 1))

    def forward(self, input3d, feature2d):
        (feature3d, output3d) = self.dense_unet_3d(input3d)
        #print(feature3d.shape, output3d.shape, feature2d.shape)
        final = torch.add(feature2d, feature3d)

        finalout = self.finalConv3d1(final)
        if (self.drop > 0):
            finalout = F.dropout(finalout, p=self.drop)

        finalout = self.finalBn(finalout)
        finalout = self.finalAc(finalout)
        finalout = self.finalConv3d2(finalout)

        return finalout

#最终的混合模型
class dense_rnn_net(nn.Module):
    def __init__(self, drop_rate=0.3):
        super(dense_rnn_net, self).__init__()
        self.drop = drop_rate
        self.dense2d = DenseUNet(in_ch=1, out_ch=3)
        self.dense3d = UNET3D(init_channels=1, class_nums=3)
        # self.conv2d5 = nn.Conv2d(64, 3, (1, 1), padding= 0)
        # self.conv3d5 = nn.Conv3d(64, 3, (1, 1, 1), padding= 0)
        self.finalConv3d1 = nn.Conv3d(64, 64, (3, 3, 3), padding=(1, 1, 1))
        self.finalBn = nn.BatchNorm3d(64)
        self.finalAc = nn.ReLU(inplace=True)
        self.finalConv3d2 = nn.Conv3d(64, 3, (1, 1, 1))

    def forward(self, data):
        input2d = data.permute(1, 0, 2, 3)
        output = self.dense2d(input2d)
        feature2d = output["feature"]

        input3d = data.unsqueeze(0)
        feature2d = feature2d.permute(1, 0, 2, 3)
        feature2d = feature2d.unsqueeze(0)

        (feature3d, output3d) = self.dense3d(input3d)

        final = torch.add(feature2d, feature3d)

        finalout = self.finalConv3d1(final)
        if (self.drop > 0):
            finalout = F.dropout(finalout, p=self.drop)

        finalout = self.finalBn(finalout)
        finalout = self.finalAc(finalout)
        finalout = self.finalConv3d2(finalout)


        return finalout
        # return output3d


if __name__ == '__main__':
    from torchsummary import summary
    unet3d = UNET3D(init_channels=1, class_nums=3).cuda()
    summary(unet3d, (1,8,512,512))